import os
import numpy as np


def make_dataset(n_samples: int, T: int, N: int, Cin: int, Cout: int):
    x = np.zeros((n_samples, T, N, Cin), dtype=np.float32)
    y = np.random.randn(n_samples, T, N, Cout).astype(np.float32)

    # channel 0: continuous feature
    x[..., 0] = np.random.rand(n_samples, T, N).astype(np.float32)

    # channels 1..11: categorical + flags + weather-like (total Cin=12)
    # 1: slot [0..95]
    x[..., 1] = np.random.randint(0, 96, (n_samples, T, N))
    # 2: day [0..6]
    x[..., 2] = np.random.randint(0, 7, (n_samples, T, N))
    # 3..6: four binary flags
    x[..., 3] = np.random.randint(0, 2, (n_samples, T, N))
    x[..., 4] = np.random.randint(0, 2, (n_samples, T, N))
    x[..., 5] = np.random.randint(0, 2, (n_samples, T, N))
    x[..., 6] = np.random.randint(0, 2, (n_samples, T, N))
    # 7: util [0..9]
    x[..., 7] = np.random.randint(0, 10, (n_samples, T, N))
    # 8: plan [0..35]
    x[..., 8] = np.random.randint(0, 36, (n_samples, T, N))
    # 9..11: three weather-like floats
    x[..., 9] = np.random.randn(n_samples, T, N).astype(np.float32)
    x[..., 10] = np.random.randn(n_samples, T, N).astype(np.float32)
    x[..., 11] = np.random.randn(n_samples, T, N).astype(np.float32)

    return x, y


def main():
    os.makedirs('data/SINPA', exist_ok=True)
    T = 12
    N = 8
    Cin = 12
    Cout = 1

    for name, n in [('train', 64), ('val', 16), ('test', 16)]:
        x, y = make_dataset(n, T, N, Cin, Cout)
        np.savez(os.path.join('data', 'SINPA', name + '.npz'), x=x, y=y)
        print(f'Wrote {name}.npz: x={x.shape}, y={y.shape}')


if __name__ == '__main__':
    main()